import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import math
from scipy.special import erfcinv
import scipy.io as scio

CONST_16QAM = np.array([
    -3 + 1j * 3,
    -3 + 1j * 1,
    -3 + 1j * -3,
    -3 + 1j * -1,
    -1 + 1j * 3,
    -1 + 1j * 1,
    -1 + 1j * -3,
    -1 + 1j * -1,
    3 + 1j * 3,
    3 + 1j * 1,
    3 + 1j * -3,
    3 + 1j * -1,
    1 + 1j * 3,
    1 + 1j * 1,
    1 + 1j * -3,
    1 + 1j * -1,
] / np.sqrt(10))

CONST_QPSK = np.array([
    -1 + 1j,
    -1 - 1j,
    1 + 1j,
    1 - 1j,
]) / np.sqrt(2)






def ofdm_modulate(prbs, num_sub_carriers, const=CONST_QPSK):
    from math import log2
    # shape of prbs: [*, bits]
    ''' S/P '''
    prbs_group = prbs.copy()
    ceil_div = lambda y, x: y // x + (y % x != 0)
    prbs_len = prbs_group.shape[-1]
    frame_len = ceil_div(prbs_len, num_sub_carriers)
    if frame_len * num_sub_carriers != prbs_len:
        prbs_cat = np.random.randint(
            0,
            2,
            size=(*prbs_group.shape[0:-1],
                  frame_len * num_sub_carriers - prbs_len))
        prbs_group = np.concatenate([prbs_group, prbs_cat], axis=-1)
    prbs_group = prbs_group.reshape(*prbs_group.shape[0:-1], num_sub_carriers,
                                    -1)
    ''' end for S/P '''

    # shape of prbs_group: [*, num_sub_carriers, bits]
    symbols = map(prbs_group,
                  constellations=const,
                  mod_order=int(log2(const.shape[0])))
    num_symbols = symbols.shape[-1]

    symbols_ofdm = ifft(symbols, axis=-1)

    symbols_ofdm = symbols_ofdm.reshape(*symbols_ofdm.shape[0:-2], -1)
    return symbols_ofdm


def ofdm_demodulate(symbols, num_sub_carriers, const=CONST_QPSK):
    from math import log2
    symbols = symbols.reshape(*symbols.shape[0:-1], num_sub_carriers, -1)
    symbols = fft(symbols, axis=-1)
    symbols = symbols.reshape(*symbols.shape[0:-2], -1)

    prbs = demap(symbols, const, mod_order=int(log2(const.shape[0])))
    return prbs


def dBm2w(x):
    return 10**(x / 10 - 3)


def colorfulPlot(seq,
                 prbs,
                 mod_order=4,
                 show=False,
                 saveFig=False,
                 savePath=None):
    seq = seq.squeeze()
    prbs = prbs.squeeze()
    assert (len(seq) == len(prbs) / mod_order)
    if saveFig == True:
        if savePath is None:
            raise Exception(
                'savePath is supposed to be assigned when saveFig is True.')
    markercolors = [(151 / 255, 173 / 255, 172 / 255, 178 / 255),
                    (173 / 255, 137 / 255, 118 / 255, 178 / 255),
                    (0, 34 / 255, 40 / 255, 178 / 255),
                    (255 / 255, 252 / 255, 153 / 255, 178 / 255),
                    (38 / 255, 188 / 255, 213 / 255, 178 / 255),
                    (167 / 255, 220 / 255, 224 / 255, 178 / 255),
                    (225 / 255, 233 / 255, 220 / 255, 178 / 255),
                    (28 / 255, 120 / 255, 135 / 255, 178 / 255),
                    (174 / 255, 221 / 255, 129 / 255, 178 / 255),
                    (107 / 255, 194 / 255, 53 / 255, 178 / 255),
                    (6 / 255, 128 / 255, 67 / 255, 178 / 255),
                    (38 / 255, 157 / 255, 128 / 255, 178 / 255),
                    (205 / 255, 179 / 255, 128 / 255, 178 / 255),
                    (3 / 255, 101 / 255, 100 / 255, 178 / 255),
                    (3 / 255, 22 / 255, 52 / 255, 178 / 255),
                    (3 / 255, 54 / 255, 73 / 255, 178 / 255)]

    clr = [0] * len(seq)
    for i in range(len(seq)):
        which_color = 0
        for m in range(mod_order):
            which_color = which_color + prbs[i * mod_order +
                                             m] * (2**(mod_order - m - 1))
        clr[i] = markercolors[which_color]

    fig = plt.scatter(seq.real, seq.imag, color=clr, s=2)

    if saveFig:
        plt.savefig(savePath)

    if show:
        plt.show()

    plt.close()
    return fig


def plot_eyediagram(samples,
                    period,
                    memory,
                    time_shift,
                    show=False,
                    upsample_factor=1):
    # Batch mode is NOT supported
    # validate legality of input

    if upsample_factor > 1:
        from scipy.signal import resample
        samples = resample(samples, num=samples.shape[-1] * upsample_factor)
    if not samples.shape[-1] % period == 0:
        raise Exception("The number of samples should be divisible by period.")
    if not time_shift < period:
        raise Exception("time_shift should be less than period.")

    x = np.linspace(1, period, period)
    # x=np.repeat(x,int(samples.shape[0]/period))
    if np.issubdtype(samples.dtype, np.complex):
        Warning("The samples are complex, imaginary part is ignored.")
    samples = samples.real

    fig = None
    for indx in range(memory):
        fig = plt.plot(
            x, samples[..., time_shift + indx * period:(indx + 1) * period +
                       time_shift])
    # fig=plt.scatter(x,samples)
    if show:
        plt.show()
    return fig


def plot_fig(fig):
    plt.show(fig)


def moving_average(x: np.ndarray, win_size: int) -> np.ndarray:
    # shape of x: [*, samples]
    input_is_complex = True in np.iscomplex(x)
    block_size = x.shape[-1]
    h = np.ones((*[1] * (x.ndim - 1), win_size)) / win_size

    x = np.fft.fft(x)
    x = x * np.fft.fft(h, n=block_size)
    x = np.fft.ifft(x)
    x = np.roll(x, -round(math.floor((win_size) // 2)), axis=-1)

    if not input_is_complex:
        x = x.real
    return x


def moving_variance(x: np.ndarray, win_size: int) -> np.ndarray:
    x_moving_average = moving_average(x, win_size)
    bs = BlockSelector(x.shape[-1], block_size=win_size, step=1)
    ret = np.zeros_like(x)
    for indx in range(len(bs)):
        res_power = np.abs(x[..., bs[indx]] -
                           x_moving_average[..., np.newaxis, indx])**2
        ret[..., indx] = np.mean(res_power, axis=-1, keepdims=False)

    return ret


def phaseRec(sig, win_size):
    def movAvr(x, win_size):
        pol = x.shape[-2]
        block_size = x.shape[-1]
        h = np.ones((pol, win_size)) / win_size
        x = np.fft.fft(x)
        x = x * np.fft.fft(h, n=block_size)
        x = np.fft.ifft(x)
        x = np.roll(x, -round(math.floor((win_size) // 2)), axis=-1)
        return x

    sig_4th_power = sig * sig * sig * sig
    sig_4th_power = movAvr(sig_4th_power * np.exp(1j * np.pi), win_size)
    phi = np.unwrap(np.angle(sig_4th_power)) / 4
    return sig * np.exp(-1j * phi)


def cal_ber(prbs, prbs_ori):
    assert (np.array_equal(prbs.shape, prbs_ori.shape))
    prbs_len = prbs.shape[-1]
    err_num = np.sum(prbs != prbs_ori, axis=-1)
    ber = err_num / prbs_len
    return ber


def decide_indx(symbols, constellations=CONST_16QAM, mod_order=4):
    ''' symbol.shape = [*, samples] '''
    constellation_points = constellations
    constellation_points = constellation_points.reshape(
        *[1] * symbols.ndim, -1)
    symbols = symbols / (np.sqrt(
        np.mean(symbols * symbols.conj(), axis=-1, keepdims=True)) + 1e-16)
    symbol_dec = np.argmin(np.abs(symbols[..., np.newaxis] -
                                  constellation_points),
                           axis=-1).astype(np.uint8)
    return symbol_dec


def decide(symbols, constellations=CONST_16QAM, mod_order=4):
    """[summary]
        Decide symbols by the nearest ideal constellation point.
    Parameters
    ----------
    symbols : [type] np.ndarray
        [description] symbols to be decided, with the shape
        of [*, no. of symbols]
    constellations : [type], optional
        [description], ideal constellation point, by default 
        CONST_16QAM
    mod_order : int, optional
        [description], by default 4

    Returns
    -------
    [type] np.ndarray
        [description] the vector of corresponding nearest ideal
        constellation points.
    """
    indxes = decide_indx(symbols, constellations, mod_order)
    symbolDec = constellations[indxes]
    return symbolDec


def demap(symbols, constellations=CONST_16QAM, mod_order=4, order='msb'):
    ''' symbol.shape = [*, samples] '''
    ''' Only for 16QAM '''
    # constellation_points = constellations
    # constellation_points = constellation_points.reshape(*[1]*symbols.ndim, -1)
    # symbols=symbols/(np.sqrt(np.mean(symbols*symbols.conj(), axis=-1, keepdims=True))+1e-16)
    # symbol_dec = np.argmin(np.abs(
    #     symbols[..., np.newaxis]-constellation_points), axis=-1).astype(np.uint8)
    # ret = np.flip(np.unpackbits(
    #     symbol_dec[..., np.newaxis], count=-mod_order, axis=-1, bitorder='little'), axis=-1)
    symbol_dec = decide_indx(symbols, constellations, mod_order)
    bits = np.unpackbits(symbol_dec[..., np.newaxis],
                         axis=-1,
                         bitorder='little')

    if order == 'msb':
        ret = np.flip(bits[..., :mod_order], axis=-1)
    elif order == 'lsb':
        ret = bits[..., :mod_order]
    else:
        raise Exception('Invalid order, which is supposed to be msb or lsb.')
    ret = ret.reshape(*ret.shape[0:-2], ret.shape[-1] * ret.shape[-2])
    return ret


def map(prbs, constellations=CONST_16QAM, mod_order=4, order='msb'):

    base_vec = np.arange(mod_order)
    base_vec = 2**base_vec
    if order == 'msb':
        base_vec = np.flip(base_vec)
    elif order == 'lsb':
        pass
    else:
        raise Exception('Invalid bit order.')

    # prbs_len = prbs.shape[-1]
    # prbs_reshape = prbs.reshape()

    symbol_dec = np.sum(prbs.reshape(*prbs.shape[0:-1], -1, mod_order) *
                        base_vec.reshape(*[1] * prbs.ndim, -1),
                        axis=-1)
    ret = constellations[symbol_dec]
    return ret


def pr_ber(sig, prbs, constellations=CONST_16QAM, mod_order=None, order='msb'):
    ''' sig.shape = [*, smpl_num] '''

    # if mod_order != 4:
    #     raise NotImplementedError

    # if mod_order is not None:

    if sig.ndim == 1:
        sig = sig.reshape(1, -1)

    if prbs.ndim == 1:
        prbs = prbs.reshape(1, -1)

    mod_order = int(math.log2(len(constellations)))

    assert (len(constellations) == 2**mod_order)
    assert (sig.shape[-1] * mod_order == prbs.shape[-1])
    sigpr = phaseRec(sig, sig.shape[-1])

    # Sub mean
    sigpr = sigpr - np.mean(sigpr, axis=-1, keepdims=True)

    rotate_fac = np.exp(1j * np.array([0, np.pi / 2, np.pi, np.pi / 2 * 3]))
    rotate_fac = rotate_fac.reshape(*[1] * sig.ndim, -1)
    sigpr = sigpr[..., np.newaxis]
    sigpr = sigpr * rotate_fac
    sigpr = np.swapaxes(sigpr, -1, -2)
    prbs = prbs[..., np.newaxis]
    prbs = np.concatenate([prbs] * 4, axis=-1)
    prbs = np.swapaxes(prbs, -1, -2)
    ''' prbs & prbs_sig & sigpr shape: [*, 4 phases, bits or samples]'''
    prbs_sig = demap(sigpr, constellations, mod_order, order)
    ''' ber shape:[*, 4 phases] '''
    ber = cal_ber(prbs, prbs_sig)
    ber_opt = np.min(ber, axis=-1)
    good_rotate_indx = np.argmin(ber, axis=-1).squeeze()
    good_rotate_indx = good_rotate_indx.reshape(
        *good_rotate_indx.shape, *[1] * (sigpr.ndim - good_rotate_indx.ndim))
    sigpr = np.swapaxes(sigpr, -1, -2)
    ret = np.take_along_axis(sigpr, good_rotate_indx, axis=-1).squeeze()
    good_rotate = rotate_fac.squeeze()[good_rotate_indx.squeeze()]
    return ret, ber_opt, good_rotate


def ber2q(ber):
    ret = 20 * np.log10(np.sqrt(2) * erfcinv(2 * ber))
    return ret

def ber2Q2(ber):
    ret = 20 * np.log10(np.sqrt(10) * erfcinv(8 * ber / 3))
    return ret

def effective_snr(sig,
                  prbs,
                  constellations=CONST_16QAM,
                  mod_order=4,
                  dB_out=True,
                  power_norm=True,
                  order='msb',
                  recover_phase=True):
    '''
        sig.shape = [*, samples]
        prbs.shape = [*, bits]
    '''

    sig_ideal = map(prbs, constellations, mod_order, order).squeeze()
    if sig_ideal.ndim == 1 and sig.ndim == 2:
        sig_ideal = sig_ideal.reshape(1, -1)

    if not np.array_equal(sig_ideal.shape, sig.shape):
        raise Exception('Shape mismatch: prbs and sig')

    if recover_phase:
        sig, ber, phase_state = pr_ber(sig,
                                    prbs,
                                    constellations=constellations,
                                    mod_order=mod_order,
                                    order=order)

    sig = sig - np.mean(sig, axis=-1, keepdims=True)
    if power_norm:
        sig = sig / np.sqrt(np.mean(abs(sig)**2, axis=-1, keepdims=True))

    ret = np.sum(np.abs(sig_ideal)**2, axis=-1) / np.sum(
        np.abs(sig - sig_ideal)**2, axis=-1)
    if dB_out == True:
        ret = 10 * np.log10(ret)
    return ret


def sym_decide_np(symbols, constellations):
    """[for numpy]
    Make decisions of the input symbols. Pick the nearest ideal symbol
    according to the value of input symbols.

    Parameters
    ----------
    symbols : numpy array
        Input symbols array.
    constellations : numpy array
        Constellations of the ideal symbols.

    Returns
    -------
    numpy array
        Symbol decisions.
    """
    const = constellations
    const = const.reshape(*[1] * symbols.ndim, -1)
    symbols = symbols / \
        (np.sqrt(np.mean(symbols*symbols.conj(), axis=-1, keepdims=True)) + 1e-20)
    symbol_dec_indx = np.argmin(np.abs(symbols[..., np.newaxis] - const),
                                axis=-1).astype(np.uint8)

    symbol_dec = np.take(const, symbol_dec_indx)

    return symbol_dec


def snr_dd(
    sig,
    constellations=CONST_16QAM,
    dB_out=True,
    power_norm=True,
):
    sig_ideal = sym_decide_np(sig, constellations)
    if power_norm:
        sig = sig / np.sqrt(np.mean(abs(sig)**2, axis=-1, keepdims=True))
    ret = np.sum(np.abs(sig_ideal)**2, axis=-1) / np.sum(
        np.abs(sig - sig_ideal)**2, axis=-1)
    if dB_out == True:
        ret = 10 * np.log10(ret)
    return ret


def cal_tap_num(DL, sample_freq, ref_freq=193.1e12):
    # [1] A. Sheikh, C. Fougstedt, A. G. I. Amat, P. Johannisson, P. Larsson-Edefors, and M. Karlsson, “Dispersion Compensation FIR Filter with Improved Robustness to Coefficient Quantization Errors,” J. Light. Technol., vol. 34, no. 22, pp. 5110–5117, 2016.
    c = 3e8
    lmbda = c / ref_freq
    taps = DL * lmbda**2 / (c * (1 / sample_freq)**2)
    return taps


def cal_tap_num_tst(D, L, df, sample_freq, ref_freq=193.1e12):

    c = 3e8
    lmbda = c / ref_freq
    beta2 = -(lmbda**2 / (2 * np.pi * c)) * D
    taps = 2 * np.pi * np.abs(beta2) * df * sample_freq * L
    return taps


def beta2toD(beta2, ref_freq=193.1e12):
    c = 3e8
    lmbda = c / ref_freq
    D = -beta2 / (lmbda**2 / (2 * np.pi * c))
    return D


def n2togamma(n2, ref_freq=193.1e12, Aeff=80.0e-12):
    omg0 = 2 * np.pi * ref_freq
    c = 3e8
    gm = omg0 * n2 / c / Aeff
    return gm


def db(x):
    ret = 10 * np.log10(x)
    return ret


class BlockSelector():
    '''
        Block selector for selecting sample blocks with overlap.

        Sketch map:
             |===============================......==========|
             |<--------------- sample_num ------------------>|
        
        |<------------- block_size------------->|
        |====|=============================|====|                                  <---Block 1
        |    | <-------   step   --------> |    |
          |                                   |
          V                                   V
      pre_overhead                        post_overhead

                                      |====|=============================|====|    <---Block 2
                                      |    | <-------   step   --------> |    |
                                         |                                  |
                                         V                                  V
                                    pre_overhead                        post_overhead
    '''
    def __init__(self, sample_num, block_size, step):
        from math import ceil, floor
        self.sample_num = sample_num
        self.block_size = block_size
        self.step = step
        self.len = round(ceil(sample_num / step))
        self.pre_overhead = round(ceil((block_size - step) / 2))
        self.post_overhead = round(floor((block_size - step) / 2))
        assert (self.pre_overhead + self.post_overhead +
                self.step == self.block_size)

    def __len__(self):
        return self.len

    def __getitem__(self, indx):
        wrapper = lambda x, N: x % N
        if isinstance(indx, np.ndarray):
            assert (indx.ndim == 1)
            raise NotImplementedError

        else:
            head = indx * self.step
            block_start = head - self.pre_overhead
            block_end = head + self.step + self.post_overhead
            pickIndx = np.arange(block_start, block_end)
            pickIndx = wrapper(pickIndx, self.sample_num)
            return pickIndx


if __name__ == '__main__':
    # ''' map demap test'''
    prbs_len = 8192
    prbs = np.random.randint(low=0, high=2, size=(2, prbs_len))

    mod_order = 4
    order = 'lsb'
    # mod_order=2
    # constellations= np.array([-1+1j,-1-1j,1+1j,1-1j,])/np.sqrt(2)
    sym = map(prbs, CONST_16QAM, mod_order, order=order)
    ''' test decide'''
    sym = decide(sym, CONST_16QAM, mod_order)

    prbs_demap = demap(sym, CONST_16QAM, mod_order, order=order)

    ber = cal_ber(prbs, prbs_demap)
    print(ber)
    # ''' end '''
    ''' cal tap number '''
    # D = 17e-6
    # # one span
    # L = 3200e3
    # symbol_rate = 32e9
    # sample_freq = 64e9
    # # beta2 = -21.668e-27
    # tap = cal_tap_num_tst(17e-6, L, df=symbol_rate, sample_freq=sample_freq)
    # print(tap)
    ''' end for cal tap number '''
